from torchvision import transforms
from copy import deepcopy

from torch.utils.data import DataLoader

from dataset.autoaugment import CIFAR10Policy
from dataset.cifar import CIFAR10, CIFAR100


class CifarDataloader:

    def __init__(self, args=None):
        self.transform_train = transforms.Compose([
            transforms.RandomCrop(32, padding=4),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
        ])

        self.transform_val = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
        ])
        self.transform_st = transforms.Compose(
            [
                transforms.RandomCrop(32, padding=4),
                transforms.RandomHorizontalFlip(),
                CIFAR10Policy(),
                transforms.ToTensor(),
                transforms.Normalize((0.507, 0.487, 0.441), (0.267, 0.256, 0.276)),
            ]
        )
        self.args = args
        self.sampler = None
        self.dataset = None
        self.warmup_loader = None

    def test(self):
        args = self.args
        if args.dataset == 100:
            dataset = CIFAR100(train=False, transform=self.transform_val)
        elif args.dataset == 10:
            dataset = CIFAR10(train=False, transform=self.transform_val)
        dataloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=False,
                                num_workers=8, pin_memory=True)
        return dataloader

    def warmup(self):
        if self.warmup_loader is None:
            args = self.args
            if args.dataset == 100:
                dataset = CIFAR100(train=True, transform=self.transform_train, asym=args.asym,
                                r_ood=args.r_ood, r_id=args.r_id, r_imb=args.r_imb, seed=args.seed)
            elif args.dataset == 10:
                dataset = CIFAR10(train=True, transform=self.transform_train, asym=args.asym,
                                r_ood=args.r_ood, r_id=args.r_id, r_imb=args.r_imb, seed=args.seed)
            self.warmup_loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True,
                                    num_workers=8, pin_memory=True)
            self.dataset = dataset
            self.dataset.transform_st = self.transform_st
        return self.warmup_loader

    def labeled(self, pred, prob):
        args = self.args
        dataset = deepcopy(self.dataset)
        dataset.filter_labeled(pred, prob)
        dataloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True,
                                num_workers=8, pin_memory=True)
        return dataloader
    
    def unlabeled(self,pred):
        args = self.args
        dataset = deepcopy(self.dataset)
        dataset.filter_unlabeled(pred)
        dataloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True,
                                num_workers=8, pin_memory=True)

        return dataloader

    def train(self, pred, prob):
        labeled_dataloader = self.labeled(pred, prob)
        unlabeled_dataloader = self.unlabeled(pred)
        return labeled_dataloader, unlabeled_dataloader
